#!/opt/local/bin/python

from analhighbeta import linacreflected
from scipy.optimize import fminbound, fmin_tnc
from rfvectors import cavity
import csv
from numpy import nan, array
import matplotlib.pyplot as plt

def spokereflected(splitinds, Ql, filename='spokenewshort.csv'):
    filehandles = []
    data = csv.reader(open(filename, 'rU'), dialect=csv.excel_tab, delimiter=',')
    rowcounter = 0
    for row in data:
        rowcounter += 1
        if rowcounter<splitinds[0]:
            filestuff(filehandles, row, 0)
        elif rowcounter<splitinds[1]:
            filestuff(filehandles, row, 1)
        else:
            filestuff(filehandles, row, 2)

    for i in filehandles:
        i.close()

    power = []
    for i in xrange(len(splitinds)+1):
        fname = 'section'+str(i)+'.csv'
        spokeargs  = (fname,  352.21e6, nan, 50e-3, 50e-3, 500.0, 0.46, 'Ezvszspoke.dat', False)
        Qlopt = fminbound(linacreflected, 1e4, 1e7, args=spokeargs)
        print "Qlopt = %e" % Qlopt
        power.append(linacreflected(Ql=Qlopt, filename=fname, f=352.21e6, delF=nan, 
            I_b0=50e-3, I_b=50e-3, RoverQ=500.0, beta_0=0.46, Eampfile='Ezvszspoke.dat'))
    return sum(power)

def filestuff(fhandlist, dat, indnum):
    try:
        fhandlist[indnum].write(','.join(dat)+'\n')
    except IndexError:
        fhandlist.append(open('section'+str(indnum)+'.csv', 'w'))
        fhandlist[indnum].write(','.join(dat)+'\n')

def minhelper(x0):
    print "%0.0f, %0.0f" % (x0[0],x0[0]+x0[1])
    return spokereflected(splitinds=[x0[0], x0[0]+x0[1]], Ql=[1e5,1e5,1e5])

if __name__=="__main__":
    #for start in xrange(2,15):
    #    for ran in xrange(2,15):
    #        totref = spokereflected(splitinds=[start,start+ran], Ql=[2e5,2e5,2e5])
    #        print "Start = %d :: Ran %d :: Total reflected = %f kW" % (start, ran, totref/1e3)
    start, ran = 4, 5
    totref = spokereflected(splitinds=[start,start+ran], Ql=[2e5,2e5,2e5])
    print "Start = %d :: Ran %d :: Total reflected = %f kW" % (start, ran, totref/1e3)

    #totref = spokereflected(splitinds=[6,16], Ql=[2e5,2e5,2e5])
    #print "Total reflected = %f" % totref

    #res, nfeval, rc = fmin_tnc(minhelper, x0=[3,3], approx_grad=True, epsilon=2.1,
    #        bounds=[(2,15), (2,20)],
    #        maxfun=1000)
    #print "nfeval = %d :: return code = %d" % (nfeval, rc)
    #print "split locations = %0.0f, %0.0f" % (36+res[0], 36+res[0]+res[1])
    #print "Ql = (%0.3e, %0.3e, %0.3e)" % (res[2], res[3], res[4])
    #print "start = %f" % start
    #print "Total reflected after optimisation = %0.3f kW" % (minhelper(res)/1e3)

    #refl = []
    #inputrange = range(2,20)
    #for start in inputrange:
    #    res, nfeval, rc = fmin_tnc(minhelper, x0=[start,3], approx_grad=True, epsilon=2.1,
    #            bounds=[(2,20), (2,20)],
    #            maxfun=1000)
    #    print "nfeval = %d :: return code = %d" % (nfeval, rc)
    #    print "split locations = %0.0f, %0.0f" % (36+res[0], 36+res[0]+res[1])
    #    print "Ql = (%0.3e, %0.3e, %0.3e)" % (res[2], res[3], res[4])
    #    print "start = %f" % start
    #    print "Total reflected after optimisation = %0.3f kW" % (minhelper(res)/1e3)
    #    refl.append(minhelper(res))

    #fig = plt.figure(1)
    #ax = fig.add_subplot(111)
    #ax.plot(inputrange, array(refl), '-x')
    #ax.grid()
    #plt.show()

